import os
import sys
import re
import glob
import queue

import originpro as op

import numpy as np
import pandas as pd

#import matplotlib.pyplot as plt

import tkinter as tk
from tkinter import filedialog, messagebox
import threading

# 全局变量
working_folder = os.getcwd()
template_folder = os.path.join(os.path.abspath('template'))
template_path1 = os.path.join(template_folder, 'IdVg_uniformity.otpu')
template_path2 = os.path.join(template_folder, 'IdVg_diff-Lch.otpu')
template_path3 = os.path.join(template_folder, 'IdVd.otpu')
#template_path4 = os.path.join(template_folder, 'IdVg_diff-T.otpu')
csv_folder = ''
template_paths = {
    "IdVg_uniformity": template_path1,
    "IdVg_diff-Lch": template_path2,
    "IdVd": template_path3,
    #"IdVg_diff-T": template_path4
}

class RedirectText(object):
    def __init__(self, widget):
        self.widget = widget

    def write(self, string):
        self.widget.insert(tk.END, string)
        self.widget.see(tk.END)

    def flush(self):
        pass

def origin_shutdown_exception_hook(exctype, value, traceback):
    op.exit()
    sys.__excepthook__(exctype, value, traceback)

############## IdVg_uniformity_process start ##############
def IdVg_uniformity_process_region(file_prefixes, region_name):
    color_index = ['black', 'red', 'green', 'blue']
    #cmap = plt.get_cmap('viridis')
    #color_index = cmap(np.linspace(0, 1, len(file_prefixes)))

    selected_files = []
    for prefix in file_prefixes:
        pattern = os.path.join(csv_folder, f'{prefix}*')
        selected_files.extend(glob.glob(pattern))
    index = 0

    if not selected_files:
        print(f'[WARNING]No files found for {region_name}')
        return 0
    else:
        print(f'>>>Found {len(selected_files)} IdVg files with {region_name}')

    for file in selected_files:
        print(f'***Processing file: {file}')
        df = pd.read_csv(file)

        Vgs = df['Vgs']
        Ids = np.abs(df['Ids'])
        Vds = df['Vds']

        wks = op.new_sheet('w', f'{os.path.basename(file)}')
        wks.from_list(0, Vgs, lname='Vgs')
        wks.from_list(1, Ids, lname='Ids')

        if index == 0:
            gr = op.new_graph(template=template_path1)
            if gr is None:
                print(f'[ERROR]Failed to create graph with template: {template_path1}')
                return 0
            gr.name = f'G_uniformity_{region_name}'
            p1 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p1.color = color_index[index]
            p2 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p2.color = color_index[index]
        
        if index == 1:
            p3 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p3.color = color_index[index]
            p4 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p4.color = color_index[index]

        if index == 2:
            p5 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p5.color = color_index[index]
            p6 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p6.color = color_index[index]
        
        if index == 3:
            p7 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p7.color = color_index[index]
            p8 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p8.color = color_index[index]
        
        index += 1
    return 1

def IdVg_uniformity_process():
    # 选取数据
    file_prefix_regions = [
        ['IdVg_R1C1', 'IdVg_R1C3', 'IdVg_R3C1', 'IdVg_R3C3'],
        ['IdVg_R1C2', 'IdVg_R1C4', 'IdVg_R3C2', 'IdVg_R3C4'],
        ['IdVg_R2C1', 'IdVg_R2C3', 'IdVg_R4C1', 'IdVg_R4C3'],
        ['IdVg_R2C2', 'IdVg_R2C4', 'IdVg_R4C2', 'IdVg_R4C4']
    ]
    for i, file_prefix_region in enumerate(file_prefix_regions):
        if IdVg_uniformity_process_region(file_prefix_region, f'{(i + 2) * 100}L'):
            print(f'[OK]Processing {(i + 2) * 100}L done!')
############## IdVg_uniformity_process end ##############

############## IdVg_diff-Lch_process start ##############
def IdVg_diff_Lch_process_region(file_prefixes, region_name):
    color_index = ['black', 'red', 'green', 'blue']

    selected_files = []
    for prefix in file_prefixes:
        pattern = os.path.join(csv_folder, f'{prefix}*')
        selected_files.extend(glob.glob(pattern))
    index = 0

    if not selected_files:
        print(f'[WARNING]No files found for region: {region_name}')
        return 0
    else:
        print(f'>>>Found {len(selected_files)} IdVg files in region: {region_name}')

    for file in selected_files:
        print(f'***Processing file: {file}')
        df = pd.read_csv(file)

        Vgs = df['Vgs']
        Ids = np.abs(df['Ids'])

        wks = op.new_sheet('w', f'{os.path.basename(file)}')
        wks.from_list(0, Vgs, lname='Vgs')
        wks.from_list(1, Ids, lname='Ids')

        if index == 0:
            gr = op.new_graph(template=template_path2)
            if gr is None:
                print(f'[ERROR]Failed to create graph with template: {template_path1}')
                return 0
            gr.name = f'G_IdVg_diff-Lch_{region_name}'
            p1 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p1.color = color_index[index]
            p2 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p2.color = color_index[index]
        
        if index == 1:
            p3 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p3.color = color_index[index]
            p4 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p4.color = color_index[index]

        if index == 2:
            p5 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p5.color = color_index[index]
            p6 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p6.color = color_index[index]
        
        if index == 3:
            p7 = gr[0].add_plot(wks, colx='Vgs', coly='Ids')
            p7.color = color_index[index]
            p8 = gr[1].add_plot(wks, colx='Vgs', coly='Ids')
            p8.color = color_index[index]
        
        index += 1
    return 1

def IdVg_diff_Lch_process():
    file_prefix_regions = [
        ['IdVg_R1C1', 'IdVg_R1C2', 'IdVg_R2C1', 'IdVg_R2C2'],
        ['IdVg_R1C3', 'IdVg_R1C4', 'IdVg_R2C3', 'IdVg_R2C4'],
        ['IdVg_R3C1', 'IdVg_R3C2', 'IdVg_R4C1', 'IdVg_R4C2'],
        ['IdVg_R3C3', 'IdVg_R3C4', 'IdVg_R4C3', 'IdVg_R4C4']
    ]

    for i, file_prefix_region in enumerate(file_prefix_regions):
        if IdVg_diff_Lch_process_region(file_prefix_region, f'Region{i+1}'):
            print(f'[OK]Processing region {i+1} done!')
############## IdVd_diff-Lch_process end ##############

############## IdVd_process start ##############
def IdVd_process():
    file_prefix ='IdVd'
    color_index = ['black', 'red', 'green', 'blue', 'cyan', 'magenta', 'yellow', 'orange', 'purple', 'brown']  # 10 colors

    selected_files = []
    pattern = os.path.join(csv_folder, f'{file_prefix}*')
    selected_files.extend(glob.glob(pattern))
    if not selected_files:
        print(f'[WARNING]No IdVd files found')
        return 0
    else:
        print(f'>>>Found {len(selected_files)} IdVd files')

    for file in selected_files:
        print(f'***Processing file: {file}')
        data = pd.read_csv(file)

        gr = op.new_graph(template=template_path3)
        if gr is None:
            print(f'[ERROR]Failed to create graph with template: {template_path1}')
            return 0
        gr.name = f'G_{os.path.basename(file)}'

        ids_columns = [col for col in data.columns if col.startswith('Ids')]
        vds_columns = 'Vds(1)'

        wks = op.new_sheet('w', f'{os.path.basename(file)}')
        wks.from_list(0, data[vds_columns], lname = 'Vds')

        for i, col in enumerate(ids_columns):
            wks.from_list(i + 1, data[col], lname = col)

        for i, col in enumerate(ids_columns):
            p = gr[0].add_plot(wks, colx='Vds', coly=col)
            p.color = color_index[i % len(color_index)]
            
    print(f'[OK]IdVd processing done!')
############## IdVd_process end ##############

def select_csv_folder():
    global csv_folder
    csv_folder = filedialog.askdirectory()
    if csv_folder:
        csv_folder_label.config(text=csv_folder)

def start_processing():
    global csv_folder, selected_type
    if not csv_folder or not selected_type.get():
        messagebox.showerror("Error", "Please select both data folder and data type.")
        return
    
    else:
        if op and op.oext:
            sys.excepthook = origin_shutdown_exception_hook

        # 打开origin界面
        if op.oext:
            op.set_show(True)

    template_path = template_paths[selected_type.get()]
    print(f"[STATE]Selected data type: {selected_type.get()}")
    print(f"[STATE]Using template: {template_path}")

    if selected_type.get() == 'IdVg_uniformity':
        IdVg_uniformity_process()

    if selected_type.get() == 'IdVg_diff-Lch':
        IdVg_diff_Lch_process()

    if selected_type.get() == 'IdVd':
        IdVd_process()

def close_origin():
    if op.oext:
        op.exit()

def close_application():
    close_origin()
    root.destroy()

# 创建主窗口
root = tk.Tk()
root.title("Transistor Data Plotter")

# 创建并放置组件
tk.Label(root, text="Select Data Folder:").grid(row=0, column=0, padx=10, pady=10)
csv_folder_label = tk.Label(root, text="")
csv_folder_label.grid(row=0, column=1, padx=10, pady=10)
tk.Button(root, text="Browse", command=select_csv_folder).grid(row=0, column=2, padx=10, pady=10)

selected_type = tk.StringVar()
tk.Label(root, text="Select Process Type:").grid(row=1, column=0, padx=10, pady=10)
data_type_menu = tk.OptionMenu(root, selected_type, *template_paths.keys())
data_type_menu.grid(row=1, column=1, padx=10, pady=10)

tk.Button(root, text="Start Processing", command=start_processing).grid(row=2, column=0, columnspan=3, padx=10, pady=20)
tk.Button(root, text="Close Origin", command=close_origin).grid(row=3, column=0, columnspan=3, padx=10, pady=10)

# 创建文本框来显示命令行输出
output_text = tk.Text(root, wrap='word', height=10, width=100)
output_text.grid(row=4, column=0, columnspan=3, padx=10, pady=10)

# 重定向 stdout 和 stderr
sys.stdout = RedirectText(output_text)
sys.stderr = RedirectText(output_text)

# 绑定窗口关闭事件
root.protocol("WM_DELETE_WINDOW", close_application)

root.resizable(False, False)

# 运行主循环
root.mainloop()